#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
中央控制节点 (CenterNode)

作为整个机器人系统的"大脑"，该节点负责：
1.  **状态管理**: 统一管理机器人的核心状态，如是否正在处理任务、导航中、说话中等，确保在正确的时间执行正确的操作。
2.  **语音指令理解**: 接收来自ASR（语音识别）节点的原始文本。
3.  **LLM集成**: 将用户指令发送给大语言模型（ZhipuAI），利用其强大的自然语言理解能力，将自然语言指令解析为结构化的JSON命令。
4.  **任务分发**: 根据LLM返回的JSON命令，向各个功能节点（如运动控制、导航、机械臂）发布相应的指令。
5.  **ASR生命周期控制**: 智能地启动和停止ASR服务，例如，在机器人说话或执行任务时暂停识别，任务结束后再自动开启，以提高交互的流畅性和准确性。
"""

import rospy
import json
import re
from std_msgs.msg import String
from std_srvs.srv import SetBool
from zhipuai import ZhipuAI
import threading
import time

class CenterNode:
    """
    封装了中央控制逻辑的ROS节点类。
    """
    def __init__(self):
        """
        初始化CenterNode。
        """
        rospy.init_node('center_node', anonymous=True)
        
        # 初始化智谱AI大模型客户端
        self.client = ZhipuAI(api_key="XXXXXXXXXXXXXXXXXXX", timeout=30.0)  
        self.model = "glm-4-flash"
        
        # --- 系统核心状态变量 ---
        self.is_processing = False      # 是否正在处理LLM请求
        self.is_navigating = False      # 是否正在导航
        self.tts_is_playing = False     # TTS是否正在播报
        self.asr_is_active = False      # ASR是否正在运行
        
        # 状态锁，用于保护对上述状态变量的并发访问，确保线程安全
        self.state_lock = threading.Lock()
        
        # 对话历史，用于支持多轮对话
        self.conversation_history = []
        self.history_lock = threading.Lock()
        
        # --- 系统提示 (System Prompt) ---
        # 这是指导LLM如何响应和格式化输出的关键。它定义了机器人的能力，
        # 并强制要求LLM以特定的JSON格式返回功能性指令。
        self.system_prompt = """你现在是一个盲人引导机器人，具有运动控制/导航/机械臂夹取/场景识别的功能，请分析用户的要求，如果你认为用户的要求属于运动控制/导航/机械臂夹取/场景识别中的一类，请按照下属json格式响应，请严格按以下JSON格式响应：
{
  "content": "回复内容",  # 当用户的要求为"场景识别"时，"content"保持空字符串即可，不要回复任何内容
  "request": "运动控制/导航/机械臂夹取/场景识别",        # 仅用户的要求为"运动控制"或者"导航"或者"机械臂夹取"或者"场景识别"时需要，选择"运动控制"或者"导航"或者"机械臂夹取"或者"场景识别"四者之一输出，其他情况设为"NONE"
  "text": "前进/后退/左转/右转/停止",  # 仅"request"为运动控制时需要，选择"前进"或者"后退"或者"左转"或者"右转"或者"停止"五者之一输出，其他情况设为"NONE"
  "place": "116"  # 仅导航时需要        # 仅"request"为导航时需要，且只有两个选项：327，302，303,且只响应数字即可，其他情况设为"NONE"  
  "arm": "true"   # 仅机械臂夹取时需要，且只要"request"为机械臂夹取，"arm"即为true， 其他情况设为"NONE"    
  "visual": "true"  # 仅场景识别时需要，且只要"request"为场景识别，"visual"即为true，其他情况设为"NONE"       
}
例如：
如果用户说"请你介绍一下你自己"，你的响应应该是：
{
  "content": "你好，我是一个盲人引导机器人，能够帮助你进行运动控制、导航、机械臂夹取和场景识别。",
  "request": "NONE",
  "text": "NONE",
  "place": "NONE",
  "arm": "NONE",
  "visual": "NONE"
}
如果用户说"请你前进"，你的响应应该是：
{
  "content": "好的，我将前进",
  "request": "运动控制",
  "text": "前进",
  "place": "NONE",
  "arm": "NONE",
  "visual": "NONE"
}
如果用户说"请你去327"，你的响应应该是：
{
  "content": "好的，我将带您前往327房间。",
  "request": "导航",
  "text": "NONE",
  "place": "327",
  "arm": "NONE",
  "visual": "NONE"
}
如果用户说"请你带我去115"，你的响应应该是：
{
  "content": "好的，我将带你去115教室",
  "request": "导航",
  "text": "NONE",
  "place": "NONE",
  "arm": "NONE",
  "visual": "NONE"
}
如果用户说"请帮我夹一瓶水"，你的响应应该是：
{
  "content": "好的，我将帮你夹一瓶水",
  "request": "机械臂夹取",
  "text": "NONE",
  "place": "NONE",
  "arm": "true",
  "visual": "NONE"
}
如果用户说"请帮我描述一下眼前的景象"或者类似，你的响应应该是：
{
  "content": "",   # 不用回答，保持空字符串即可
  "request": "场景识别",
  "text": "NONE",
  "place": "NONE",
  "arm": "NONE",
  "visual": "true"
}
"""
        
        # --- ROS Subscribers ---
        # 订阅ASR节点的原始识别文本
        rospy.Subscriber('raw_text', String, self.voice_command_callback)
        # 订阅TTS节点的状态，以决定何时可以重新启动ASR
        rospy.Subscriber('tts_status', String, self.tts_status_callback)
        # 订阅导航节点的状态，以决定何时可以重新启动ASR
        rospy.Subscriber('navigation_status', String, self.nav_status_callback)
        
        # --- ROS Publishers ---
        # 用于向各功能节点分发指令
        self.command_pub = rospy.Publisher('robot_command', String, queue_size=10, latch=True) # 运动控制
        self.nav_pub = rospy.Publisher('navigation_command', String, queue_size=10)             # 导航
        self.arm_pub = rospy.Publisher('arm_command', String, queue_size=10)                    # 机械臂
        self.visual_pub = rospy.Publisher('visual_command', String, queue_size=10, latch=True)  # 视觉识别
        self.speech_pub = rospy.Publisher('tts_text', String, queue_size=10)                    # 语音合成
        
        # --- ROS Service Clients ---
        # 用于控制ASR节点的启动和停止
        rospy.loginfo("正在等待ASR服务...")
        rospy.wait_for_service('start_asr')
        rospy.wait_for_service('stop_asr')
        self.start_asr_client = rospy.ServiceProxy('start_asr', SetBool)
        self.stop_asr_client = rospy.ServiceProxy('stop_asr', SetBool)
        rospy.loginfo("ASR服务已连接。")
        
        # 初始化对话历史，并启动ASR，使机器人进入聆听状态
        self.initialize_history()
        self.start_asr()

    def initialize_history(self):
        """初始化或重置对话历史，将系统提示作为第一条消息。"""
        with self.history_lock:
            self.conversation_history = [{"role": "system", "content": self.system_prompt}]

    def add_to_history(self, role, content):
        """
        向对话历史中添加一条新消息，并管理历史记录的长度。
        
        Args:
            role (str): 消息的角色 ("user" 或 "assistant").
            content (str): 消息内容。
        """
        with self.history_lock:
            self.conversation_history.append({"role": role, "content": content})
            # 限制历史记录的长度，防止上下文过长。保留系统提示和最近的5轮对话。
            if len(self.conversation_history) > 11: # 1 (system) + 5 * 2 (user/assistant)
                self.conversation_history = [self.conversation_history[0]] + self.conversation_history[-10:]

    def nav_status_callback(self, msg):
        """
        处理来自导航节点的状态更新。
        主要用于在导航开始时停止ASR，在导航结束后尝试重启ASR。
        """
        with self.state_lock:
            was_navigating = self.is_navigating
            self.is_navigating = (msg.data == "active")
            
            if self.is_navigating and not was_navigating:
                rospy.loginfo("导航任务开始，暂停ASR以避免干扰。")
                self._stop_asr_internal()
            elif not self.is_navigating and was_navigating:
                rospy.loginfo("导航任务结束。")
                # 导航结束后，检查是否可以安全地重启ASR
                if not self.is_processing and not self.tts_is_playing:
                    rospy.loginfo("系统空闲，准备重新启动ASR。")
                    self._start_asr_internal()

    def start_asr(self):
        """公共接口：请求启动ASR。"""
        with self.state_lock:
            self._start_asr_internal()

    def stop_asr(self):
        """公共接口：请求停止ASR。"""
        with self.state_lock:
            self._stop_asr_internal()

    def _start_asr_internal(self):
        """
        内部实现：启动ASR。
        此方法必须在获取state_lock后调用。
        它会检查是否满足启动条件（没有其他任务正在进行）。
        """
        if not self.asr_is_active and not self.tts_is_playing and not self.is_navigating:
            try:
                # 调用服务来启动ASR节点
                # 先调用停止以确保状态一致性
                self.stop_asr_client(True)  
                rospy.sleep(0.1)
                result = self.start_asr_client(True)
                if result.success:
                    self.asr_is_active = True
                    self.is_processing = False # ASR启动后，处理流程结束
                    rospy.loginfo("ASR已成功启动。")
                else:
                    rospy.logerr(f"启动ASR服务失败: {result.message}")
            except Exception as e:
                rospy.logerr(f"调用ASR启动服务时发生异常: {e}")

    def _stop_asr_internal(self):
        """
        内部实现：停止ASR。
        此方法必须在获取state_lock后调用。
        """
        if self.asr_is_active:
            try:
                # 调用服务来停止ASR节点
                result = self.stop_asr_client(True)
                if result.success:
                    self.asr_is_active = False
                    rospy.loginfo("ASR已成功停止。")
                else:
                    rospy.logerr(f"停止ASR服务失败: {result.message}")
            except Exception as e:
                rospy.logerr(f"调用ASR停止服务时发生异常: {e}")

    def voice_command_callback(self, msg):
        """
        处理来自ASR的原始语音指令文本，是用户交互的主要入口。
        """
        with self.state_lock:
            # 如果正在处理上一个指令或正在导航，则忽略新指令
            if self.is_processing:
                rospy.logwarn("正在处理上一条指令，忽略当前输入。")
                return
            if self.is_navigating:
                rospy.logwarn("正在导航中，忽略语音指令。")
                return
                
            # 设置处理中状态，并立即停止ASR，防止进一步输入
            self.is_processing = True
            self._stop_asr_internal()
        
        user_text = msg.data.strip()
        if not user_text:
            # 如果是空文本，则直接恢复空闲状态
            with self.state_lock:
                self.is_processing = False
                self._start_asr_internal() # 尝试重启ASR
            return

        rospy.loginfo(f"收到用户指令: '{user_text}'")
        self.add_to_history("user", user_text)
        
        # 在新的后台线程中调用LLM，避免阻塞ROS回调线程
        threading.Thread(target=self.process_with_llm, daemon=True).start()

    def process_with_llm(self):
        """
        在后台线程中调用大语言模型，处理对话历史，并处理流式响应。
        """
        is_json_response = False
        full_response_content = ""
        buffer = ""
        
        try:
            with self.history_lock:
                rospy.loginfo("正在向大模型发送请求...")
                # 以流式方式调用LLM
                response_stream = self.client.chat.completions.create(
                    model=self.model,
                    messages=self.conversation_history,
                    stream=True,
                )

            # 逐块处理流式响应
            for chunk in response_stream:
                content = chunk.choices[0].delta.content
                if content is None:
                    continue

                rospy.logdebug(f"收到LLM流式数据: {content}")

                # 简单的启发式方法，检测响应是否可能是JSON
                if not full_response_content and '{' in content:
                    is_json_response = True
                
                full_response_content += content
                
                if is_json_response:
                    # 如果认为是JSON，则尝试解析
                    # 查找最后一个 '}' 来判断JSON对象是否结束
                    if '}' in content:
                        try:
                            # 提取并解析JSON
                            json_part = full_response_content[full_response_content.find('{'):full_response_content.rfind('}')+1]
                            data = json.loads(json_part)
                            rospy.loginfo(f"成功解析LLM返回的JSON: {data}")
                            # 执行指令
                            self.execute_function_commands(data)
                            # 成功处理后，将模型的完整JSON响应（可能包含自然语言部分）存入历史
                            self.add_to_history("assistant", json.dumps(data, ensure_ascii=False))
                            return # JSON处理完毕，退出函数
                        except json.JSONDecodeError:
                            # 如果解析失败，说明JSON还未完整，继续接收
                            rospy.logdebug("JSON仍在接收中...")
                else:
                    # 如果不是JSON，则认为是普通文本，按句子进行处理和播报
                    buffer += content
                    if any(p in buffer for p in ["。", "！", "？", "\n"]):
                        punc = next(p for p in ["。", "！", "？", "\n"] if p in buffer)
                        sentence, _, remainder = buffer.partition(punc)
                        to_speak = (sentence + punc).strip()
                        if to_speak:
                            self.speech_pub.publish(to_speak)
                        buffer = remainder
            
            # 处理循环结束后剩余的文本（非JSON情况）
            if not is_json_response:
                final_text = buffer.strip()
                if final_text:
                    self.speech_pub.publish(final_text)
                
                rospy.loginfo(f"LLM完整文本响应: {full_response_content}")
                self.add_to_history("assistant", full_response_content)
                # 等待TTS播放完毕后重启ASR
                self._wait_and_restart_asr()

        except Exception as e:
            rospy.logerr(f"与大模型交互时发生错误: {e}")
            self.speech_pub.publish("抱歉，我的大脑好像出了一点问题。")
            self._wait_and_restart_asr()

    def _wait_and_restart_asr(self):
        """
        等待TTS播放完成后，重新启动ASR。
        这是一个阻塞操作，应在后台线程中调用。
        """
        rospy.loginfo("等待TTS播放完成...")
        time.sleep(0.5) # 给予TTS开始播放的时间
        
        # 循环检查TTS是否仍在播放
        while True:
            with self.state_lock:
                if not self.tts_is_playing:
                    break
            time.sleep(0.2)
        
        rospy.loginfo("TTS播放完毕，准备重启ASR。")
        # 启动ASR，恢复聆听状态
        self.start_asr()

    def execute_function_commands(self, data):
        """
        根据从LLM解析出的JSON数据，执行相应的功能指令。

        Args:
            data (dict): 从LLM返回的JSON解析出的字典。
        """
        content = data.get("content", "").strip()
        request = data.get("request", "NONE")
        
        # 1. 播报回复内容（如果存在）
        if content:
            self.speech_pub.publish(content)

        # 2. 根据request类型分发指令
        if request == "运动控制":
            cmd = data.get("text", "NONE")
            if cmd != "NONE":
                rospy.loginfo(f"发布运动控制指令: {cmd}")
                self.command_pub.publish(cmd)
        elif request == "导航":
            place = data.get("place", "NONE")
            if place != "NONE":
                rospy.loginfo(f"发布导航指令，前往: {place}")
                self.nav_pub.publish(place)
        elif request == "机械臂夹取":
            cmd = data.get("arm", "NONE")
            if cmd == "true":
                rospy.loginfo("发布机械臂夹取指令")
                self.arm_pub.publish("true")
        elif request == "场景识别":
            cmd = data.get("visual", "NONE")
            if cmd == "true":
                rospy.loginfo("发布场景识别指令")
                self.visual_pub.publish("true")

        # 3. 决定后续操作
        # 如果有功能性请求（导航除外），通常是快速任务，完成后直接重启ASR
        # 导航任务的状态由nav_status_callback管理
        if request != "NONE" and request != "导航":
             # 如果有语音播报，等待播报完
            if content:
                self._wait_and_restart_asr()
            else: # 如果没有语音播报（如场景识别），直接重启ASR
                self.start_asr()
        # 如果只是普通聊天
        elif request == "NONE" and content:
            self._wait_and_restart_asr()
        # 如果没有任何操作
        else:
            self.start_asr()

    def tts_status_callback(self, msg):
        """
        处理来自TTS节点的状态更新。
        """
        with self.state_lock:
            self.tts_is_playing = (msg.data == "playing")
            rospy.logdebug(f"TTS状态更新: is_playing = {self.tts_is_playing}")
            # 当TTS停止播放，并且没有其他任务时，尝试重启ASR
            if not self.tts_is_playing and not self.is_processing and not self.is_navigating:
                self._start_asr_internal()

if __name__ == '__main__':
    try:
        node = CenterNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("中央控制节点关闭。")